package com.mori.encrypt.spring.boot.core.wrapper;

import com.mori.encrypt.spring.boot.core.util.StreamUtils;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * <p> 加密请求包装类 </p>
 *
 * @author mori
 * @since 2020/9/29 16:53
 */
public class EncryptionRequestWrapper extends HttpServletRequestWrapper {

	private byte[] requestBody;

	private Map<String, String> paramMap = new HashMap<>();

	public EncryptionRequestWrapper(HttpServletRequest request) {
		super(request);
		try {
			requestBody = StreamUtils.copyToByteArray(request.getInputStream());
		} catch (IOException e) {
			throw new RuntimeException(e);
		}
	}

	@Override
	public ServletInputStream getInputStream() {
		final ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(requestBody);
		return new ServletInputStream() {
			@Override
			public int read() {
				return byteArrayInputStream.read();
			}

			@Override
			public boolean isFinished() {
				return false;
			}

			@Override
			public boolean isReady() {
				return true;
			}

			@Override
			public void setReadListener(ReadListener listener) {

			}
		};
	}

	public String getRequestData() {
		return new String(requestBody);
	}

	public void setRequestData(String requestData) {
		this.requestBody = requestData.getBytes();
	}

	public void setParamMap(Map<String, String> paramMap) {
		this.paramMap = paramMap;
	}

	@Override
	public String getParameter(String name) {
		if (paramMap.containsKey(name)) {
			return this.paramMap.get(name);
		}
		return super.getParameter(name);
	}

	@Override
	public String[] getParameterValues(String name) {
		if (paramMap.containsKey(name)) {
			return new String[]{getParameter(name)};
		}
		return super.getParameterValues(name);
	}

	@Override
	public Map<String, String[]> getParameterMap() {
		return paramMap.entrySet()
				.stream()
				.collect(Collectors.toMap(
						Map.Entry::getKey,
						(v) -> new String[]{v.getValue()},
						(v1, v2) -> v2,
						() -> new HashMap<>(super.getParameterMap())));
	}
}
